package com.cl.spark.node;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.cl.spark.base.BaseSparkNode;
import com.cl.spark.dto.SparkParam;
import com.cl.spark.dto.SparkResult;
import com.cl.spark.enums.SparkNodeEnum;
import org.apache.spark.sql.*;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

@Component
public class FunctionNode extends BaseSparkNode {
    @Autowired
    SparkSession sparkSession;

    @Override
    public SparkResult process(SparkParam sparkParam) {
        Dataset<Row> dataset = sparkParam.getSparkResultList().get(0).getDataset();
        JSONObject expression = sparkParam.getNodeExpression();

        if (expression == null) {
            return SparkResult.success(getDatasetCount(dataset));
        }
        String function = expression.getString("function");
        if ("计数".equals(function)) {
            return SparkResult.success(getDatasetCount(dataset));
        }

        return SparkResult.success(getDatasetCount(dataset));

    }

    public Dataset<Row> getDatasetCount(Dataset<Row> dataset) {
        long count = dataset.count();
        return dataset.withColumn("count", functions.lit(count));
    }

    @Override
    public SparkNodeEnum getType() {
        return SparkNodeEnum.FUNCTION;
    }
}
